{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b8b808d5",
   "metadata": {},
   "source": [
    "# M5 Agentic AI - Customer Service Agent\n",
    "\n",
    "## 1. Introduction\n",
    "\n",
    "As Andrew explained in the lecture, *planning with code execution* means letting the LLM **write code that becomes the plan itself**.  \n",
    "Compared to plain-text or JSON-based plans, this approach is more expressive and flexible: the code not only documents the steps but can also execute them directly.\n",
    "\n",
    "In this lab, you will implement this design pattern in practice.  \n",
    "Instead of asking the LLM to output a plan in JSON format and then manually executing each step, we will allow it to **write Python code** that directly captures multiple steps of a plan. By executing this code, we can carry out complex queries automatically.  \n",
    "\n",
    "To make things concrete, we simulate a **sunglasses store** with an **inventory** of products and a set of **transactions** (sales, returns, balance updates). This example shows how the LLM can generate code to query or update records, demonstrating the flexibility of this pattern.\n",
    "\n",
    "### 1.1 Lab Overview\n",
    "We will:\n",
    "1. Create simple **inventory** and **transaction** datasets.  \n",
    "2. Build a **schema block** describing the data.  \n",
    "3. Prompt the LLM to **write a plan as Python code** (with comments explaining each step).  \n",
    "4. Execute the code in a sandbox to obtain the answer.  \n",
    "\n",
    "### 1.2 Learning Outcomes\n",
    "\n",
    "By the end of this lab, you will be able to:\n",
    "\n",
    "- **Explain** why letting the model write code (instead of JSON or plain text plans) enables richer, more flexible planning.  \n",
    "- **Prompt** an LLM to produce Python code with step-by-step comments that both documents and executes the plan.  \n",
    "- **Run** the generated code safely in a sandbox and interpret the results.  \n",
    "\n",
    "This illustrates how *Code as Action* can outperform brittle tool chains and JSON-based planning approaches."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "866b84e1",
   "metadata": {},
   "source": [
    "## 2. Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42f871b5",
   "metadata": {
    "height": 283
   },
   "outputs": [],
   "source": [
    "# ==== Imports ====\n",
    "from __future__ import annotations\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import re, io, sys, traceback, json\n",
    "from typing import Any, Dict, Optional\n",
    "from tinydb import Query, where\n",
    "\n",
    "# Utility modules\n",
    "import utils      # helper functions for prompting/printing\n",
    "import inv_utils  # functions for inventory, transactions, schema building, and TinyDB seeding\n",
    "\n",
    "load_dotenv()\n",
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1bb49e5",
   "metadata": {},
   "source": [
    "In the `inv_utils` module, we have functions like:\n",
    "\n",
    "- `create_inventory()` – builds the sunglasses inventory.  \n",
    "- `create_transactions()` – builds the initial transaction log.  \n",
    "- `seed_db()` – loads both inventory and transactions into a JSON-backed store.  \n",
    "- `build_schema_block()` – generates a schema description used in the prompt.  \n",
    "- Helpers like `get_current_balance()` and `next_transaction_id()` – let the LLM handle consistent updates across inventory and transactions.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39d609be",
   "metadata": {},
   "source": [
    "### 2.1 Create Example Tables\n",
    "\n",
    "We will now create two small tables for the sunglasses store simulation, using **[TinyDB](https://tinydb.readthedocs.io/)** — a lightweight document-oriented database written in pure Python.  \n",
    "TinyDB stores data as JSON documents and is well-suited for small applications or prototypes, since it requires no server setup and allows you to query and update data easily.\n",
    "\n",
    "The two tables are:\n",
    "\n",
    "- **`inventory_tbl`**: contains product details such as name, item ID, description, quantity in stock, and price.  \n",
    "- **`transactions_tbl`**: starts with an opening balance and will later track purchases, returns, and adjustments.  \n",
    "\n",
    "You will generate these tables using helper functions in `inv_utils`, and then preview the first few rows below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4005d476",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "db, inventory_tbl, transactions_tbl = inv_utils.seed_db()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bce9cfc7",
   "metadata": {},
   "source": [
    "Now, you can inspect the records in each table by printing them as formatted JSON:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbfb3daa",
   "metadata": {
    "height": 62
   },
   "outputs": [],
   "source": [
    "utils.print_html(json.dumps(inventory_tbl.all(), indent=2), title=\"Inventory Table\")\n",
    "utils.print_html(json.dumps(transactions_tbl.all(), indent=2), title=\"Transactions Table\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf279693",
   "metadata": {},
   "source": [
    "As you can see above, the schemas of each table are as follows:\n",
    "\n",
    "<div style=\"border:1px solid #BFDBFE; border-left:6px solid #3B82F6; background:#EFF6FF; border-radius:6px; padding:16px; font-family:system-ui,-apple-system,Segoe UI,Roboto,Ubuntu,Cantarell,Noto Sans,sans-serif; line-height:1.6; color:#1E3A8A;\">\n",
    "\n",
    "  <h4 style=\"margin-top:0; color:#1E40AF;\">Inventory Table (<code>inventory_tbl</code>)</h4>\n",
    "  <ul>\n",
    "    <li><strong>item_id</strong> (string): Unique product identifier (e.g., SG001).</li>\n",
    "    <li><strong>name</strong> (string): Style of sunglasses (e.g., Aviator, Round).</li>\n",
    "    <li><strong>description</strong> (string): Text description of the product.</li>\n",
    "    <li><strong>quantity_in_stock</strong> (int): Current stock available.</li>\n",
    "    <li><strong>price</strong> (float): Price in USD.</li>\n",
    "  </ul>\n",
    "  <h4 style=\"margin-top:1em; color:#1E40AF;\">Transactions Table (<code>transactions_tbl</code>)</h4>\n",
    "  <ul>\n",
    "    <li><strong>transaction_id</strong> (string): Unique identifier (e.g., TXN001).</li>\n",
    "    <li><strong>customer_name</strong> (string): Name of the customer, or <code>OPENING_BALANCE</code> for initial entry.</li>\n",
    "    <li><strong>transaction_summary</strong> (string): Short description of the transaction.</li>\n",
    "    <li><strong>transaction_amount</strong> (float): Amount of money for this transaction.</li>\n",
    "    <li><strong>balance_after_transaction</strong> (float): Running balance after applying the transaction.</li>\n",
    "    <li><strong>timestamp</strong> (string): ISO-8601 formatted date/time of the transaction.</li>\n",
    "  </ul>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31091813",
   "metadata": {},
   "source": [
    "## Planning with Code Execution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07a2646f",
   "metadata": {},
   "source": [
    "### 2.1. The plan\n",
    "\n",
    "Once the schema is clear, you’ll build the **prompt** that instructs the model to *plan by writing code* and then execute that code. As Andrew emphasized, the code is the plan: the model explains each step in comments, then carries it out. Your prompt below also makes the model self-decide whether the request is read-only or a state change, and it enforces safe execution (no I/O, no network, TinyDB Query only, consistent mutations).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67088dbe",
   "metadata": {
    "height": 1354
   },
   "outputs": [],
   "source": [
    "PROMPT = \"\"\"You are a senior data assistant. PLAN BY WRITING PYTHON CODE USING TINYDB.\n",
    "\n",
    "Database Schema & Samples (read-only):\n",
    "{schema_block}\n",
    "\n",
    "Execution Environment (already imported/provided):\n",
    "- Variables: db, inventory_tbl, transactions_tbl  # TinyDB Table objects\n",
    "- Helpers: get_current_balance(tbl) -> float, next_transaction_id(tbl, prefix=\"TXN\") -> str\n",
    "- Natural language: user_request: str  # the original user message\n",
    "\n",
    "PLANNING RULES (critical):\n",
    "- Derive ALL filters/parameters from user_request (shape/keywords, price ranges \"under/over/between\", stock mentions,\n",
    "  quantities, buy/return intent). Do NOT hard-code values.\n",
    "- Build TinyDB queries dynamically with Query(). If a constraint isn't in user_request, don't apply it.\n",
    "- Be conservative: if intent is ambiguous, do read-only (DRY RUN).\n",
    "\n",
    "TRANSACTION POLICY (hard):\n",
    "- Do NOT create aggregated multi-item transactions.\n",
    "- If the request contains multiple items, create a separate transaction row PER ITEM.\n",
    "- For each item:\n",
    "  - compute its own line total (unit_price * qty),\n",
    "  - insert ONE transaction with that amount,\n",
    "  - update balance sequentially (balance += line_total),\n",
    "  - update the item’s stock.\n",
    "- If any requested item lacks sufficient stock, do NOT mutate anything; reply with STATUS=\"insufficient_stock\".\n",
    "\n",
    "HUMAN RESPONSE REQUIREMENT (hard):\n",
    "- You MUST set a variable named `answer_text` (type str) with a short, customer-friendly sentence (1–2 lines).\n",
    "- This sentence is the only user-facing message. No dataframes/JSON, no boilerplate disclaimers.\n",
    "- If nothing matches, politely say so and offer a nearby alternative (closest style/price) or a next step.\n",
    "\n",
    "ACTION POLICY:\n",
    "- If the request clearly asks to change state (buy/purchase/return/restock/adjust):\n",
    "    ACTION=\"mutate\"; SHOULD_MUTATE=True; perform the change and write a matching transaction row.\n",
    "  Otherwise:\n",
    "    ACTION=\"read\"; SHOULD_MUTATE=False; simulate and explain briefly as a dry run (in logs only).\n",
    "\n",
    "FAILURE & EDGE-CASE HANDLING (must implement):\n",
    "- Do not capture outer variables in Query.test. Pass them as explicit args.\n",
    "- Always set a short `answer_text`. Also set a string `STATUS` to one of:\n",
    "  \"success\", \"no_match\", \"insufficient_stock\", \"invalid_request\", \"unsupported_intent\".\n",
    "- no_match: No items satisfy the filters → suggest the closest in style/price, or invite a different range.\n",
    "- insufficient_stock: Item found but stock < requested qty → state available qty and offer the max you can fulfill.\n",
    "- invalid_request: Unable to parse essential info (e.g., quantity for a purchase/return) → ask for the missing piece succinctly.\n",
    "- unsupported_intent: The action is outside the store’s capabilities → provide the nearest supported alternative.\n",
    "- In all cases, keep the tone helpful and concise (1–2 sentences). Put technical details (e.g., ACTION/DRY RUN) only in stdout logs.\n",
    "\n",
    "OUTPUT CONTRACT:\n",
    "- Return ONLY executable Python between these tags (no extra text):\n",
    "  <execute_python>\n",
    "  # your python\n",
    "  </execute_python>\n",
    "\n",
    "CODE CHECKLIST (follow in code):\n",
    "1) Parse intent & constraints from user_request (regex ok).\n",
    "2) Build TinyDB condition incrementally; query inventory_tbl.\n",
    "3) If mutate: validate stock, update inventory, insert a transaction (new id, amount, balance, timestamp).\n",
    "4) ALWAYS set:\n",
    "   - `answer_text` (human sentence, required),\n",
    "   - `STATUS` (see list above).\n",
    "   Also print a brief log to stdout, e.g., \"LOG: ACTION=read DRY_RUN=True STATUS=no_match\".\n",
    "5) Optional: set `answer_rows` or `answer_json` if useful, but `answer_text` is mandatory.\n",
    "\n",
    "TONE EXAMPLES (for `answer_text`):\n",
    "- success: \"Yes, we have our Classic sunglasses, a round frame, for $60.\"\n",
    "- no_match: \"We don’t have round frames under $100 in stock right now, but our Moon round frame is available at $120.\"\n",
    "- insufficient_stock: \"We only have 1 pair of Classic left; I can reserve that for you.\"\n",
    "- invalid_request: \"I can help with that—how many pairs would you like to purchase?\"\n",
    "- unsupported_intent: \"We can’t refurbish frames, but I can suggest similar new models.\"\n",
    "\n",
    "Constraints:\n",
    "- Use TinyDB Query for filtering. Standard library imports only if needed.\n",
    "- Keep code clear and commented with numbered steps.\n",
    "\n",
    "User request:\n",
    "{question}\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2da9d1f",
   "metadata": {},
   "source": [
    "### 2.2 From Prompt to Code (Planning in Code)\n",
    "\n",
    "Let’s generate code that **is the plan**.\n",
    "\n",
    "Instead of asking the model to output a plan in JSON and running it step-by-step with many tiny tools, let’s have it **write Python that encodes the whole plan** (e.g., “filter this, then compute that, then update this row”). The function `generate_llm_code`:\n",
    "\n",
    "1. **Builds a live schema** from `inventory_tbl` and `transactions_tbl` so the model sees real fields, types, and examples.\n",
    "2. **Formats the prompt** with that schema plus the user’s question.\n",
    "3. **Calls the model** to produce a **plan-with-code** response — typically an `<execute_python>...</execute_python>` block whose body contains the step-by-step logic.\n",
    "4. **Returns the full response** (including the plan and the code).  \n",
    "   *We don’t execute anything in this step.*\n",
    "\n",
    "Why this pattern? Let’s leverage Python/TinyDB as a rich toolbox the model already “knows,” so it can compose multi-step solutions directly in code instead of relying on a growing set of bespoke tools. We’ll extract and run the code in a later step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13d80308",
   "metadata": {
    "height": 555
   },
   "outputs": [],
   "source": [
    "# ---------- 1) Code generation ----------\n",
    "def generate_llm_code(\n",
    "    prompt: str,\n",
    "    *,\n",
    "    inventory_tbl,\n",
    "    transactions_tbl,\n",
    "    model: str = \"gpt-4.1-mini\",\n",
    "    temperature: float = 0.2,\n",
    ") -> str:\n",
    "    \"\"\"\n",
    "    Ask the LLM to produce a plan-with-code response.\n",
    "    Returns the FULL assistant content (including surrounding text and tags).\n",
    "    The actual code extraction happens later in execute_generated_code.\n",
    "    \"\"\"\n",
    "    schema_block = inv_utils.build_schema_block(inventory_tbl, transactions_tbl)\n",
    "    prompt = PROMPT.format(schema_block=schema_block, question=prompt)\n",
    "\n",
    "    resp = client.chat.completions.create(\n",
    "        model=model,\n",
    "        temperature=temperature,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": \"You write safe, well-commented TinyDB code to handle data questions and updates.\"\n",
    "            },\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "    )\n",
    "    content = resp.choices[0].message.content or \"\"\n",
    "    \n",
    "    return content  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e786ccaa",
   "metadata": {},
   "source": [
    "### 2.3 Try a Sample Prompt (Planning-in-Code)\n",
    "\n",
    "We’ll use the same prompt Andrew used in the lecture:\n",
    "\n",
    "> **Prompt:** “Do you have any round sunglasses in stock that are under $100?”\n",
    "\n",
    "Before generating any code, let’s manually inspect the TinyDB tables to see if there are truly *round* frames (word-only match) and what their prices look like. Run the next cell to preview the inventory and highlight items that match the word-only “round” filter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c74c6b55",
   "metadata": {
    "height": 266
   },
   "outputs": [],
   "source": [
    "Item = Query()                    # Create a Query object to reference fields (e.g., Item.name, Item.description)\n",
    "\n",
    "# Search the inventory table for documents where either the description OR the name\n",
    "# contains the word \"round\" (case-insensitive). The check is done inline:\n",
    "# - (v or \"\") ensures we handle None by converting it to an empty string\n",
    "# - .lower() normalizes case\n",
    "# - \" round \" enforces a crude word boundary (won't match \"wraparound\")\n",
    "round_sunglasses = inventory_tbl.search(\n",
    "    (Item.description.test(lambda v: \" round \" in ((v or \"\").lower()))) |\n",
    "    (Item.name.test(        lambda v: \" round \" in ((v or \"\").lower())))\n",
    ")\n",
    "\n",
    "# Render the results as formatted JSON in the notebook UI\n",
    "utils.print_html(json.dumps(round_sunglasses, indent=2), title=\"Inventory Status: Round Sunglasses\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbd78aea",
   "metadata": {},
   "source": [
    "Great — we do have round frames available. From our manual inspection, there are two round styles in stock, but only **one** is **under \\$100**. Therefore, the item that satisfies the requirement is:\n",
    "\n",
    "````python\n",
    "{\n",
    "  \"item_id\": \"SG005\",\n",
    "  \"name\": \"Classic\",\n",
    "  \"description\": \"Classic round profile with minimalist metal frames, offering a timeless and versatile style that fits both casual and formal wear.\",\n",
    "  \"quantity_in_stock\": 10,\n",
    "  \"price\": 60\n",
    "}\n",
    "````\n",
    "\n",
    "Now let’s ask the model to **generate a plan in code** that answers Andrew’s prompt (no execution yet)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f8e4978",
   "metadata": {
    "height": 266
   },
   "outputs": [],
   "source": [
    "# Andrew's prompt from the lecture\n",
    "prompt_round = \"Do you have any round sunglasses in stock that are under $100?\"\n",
    "\n",
    "# Generate the plan-as-code (FULL content; may include <execute_python> tags)\n",
    "full_content_round = generate_llm_code(\n",
    "    prompt_round,\n",
    "    inventory_tbl=inventory_tbl,\n",
    "    transactions_tbl=transactions_tbl,\n",
    "    model=\"o4-mini\",\n",
    "    temperature=1.0,\n",
    ")\n",
    "\n",
    "# Inspect the LLM’s plan + code (no execution here)\n",
    "utils.print_html(full_content_round, title=\"Plan with Code (Full Response)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c17e4ea",
   "metadata": {},
   "source": [
    "### 2.4. Define the executor function (run a given plan)\n",
    "\n",
    "Now we’ll define the function that **takes a plan produced by the model and runs it** safely:\n",
    "\n",
    "- It **accepts either** the full LLM response (with `<execute_python>…</execute_python>`) **or** raw Python code.\n",
    "- It **extracts** the executable block when needed.\n",
    "- It runs the code in a **controlled namespace** (TinyDB tables + safe helpers only).\n",
    "- It captures **stdout**, **errors**, and the model-set answer variables (`answer_text`, `answer_rows`, `answer_json`).\n",
    "- It renders **before/after** table snapshots to make side effects explicit.\n",
    "\n",
    "This is the “executor” that turns a **plan-as-code** into actions and a concise user-facing answer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6caf0073",
   "metadata": {
    "height": 1201
   },
   "outputs": [],
   "source": [
    "# --- Helper: extract code between <execute_python>...</execute_python> ---\n",
    "def _extract_execute_block(text: str) -> str:\n",
    "    \"\"\"\n",
    "    Returns the Python code inside <execute_python>...</execute_python>.\n",
    "    If no tags are found, assumes 'text' is already raw Python code.\n",
    "    \"\"\"\n",
    "    if not text:\n",
    "        raise RuntimeError(\"Empty content passed to code executor.\")\n",
    "    m = re.search(r\"<execute_python>(.*?)</execute_python>\", text, re.DOTALL | re.IGNORECASE)\n",
    "    return m.group(1).strip() if m else text.strip()\n",
    "\n",
    "\n",
    "# ---------- 2) Code execution ----------\n",
    "def execute_generated_code(\n",
    "    code_or_content: str,\n",
    "    *,\n",
    "    db,\n",
    "    inventory_tbl,\n",
    "    transactions_tbl,\n",
    "    user_request: Optional[str] = None,\n",
    ") -> Dict[str, Any]:\n",
    "    \"\"\"\n",
    "    Execute code in a controlled namespace.\n",
    "    Accepts either raw Python code OR full content with <execute_python> tags.\n",
    "    Returns minimal artifacts: stdout, error, and extracted answer.\n",
    "    \"\"\"\n",
    "    # Extract code here (now centralized)\n",
    "    code = _extract_execute_block(code_or_content)\n",
    "\n",
    "    SAFE_GLOBALS = {\n",
    "        \"Query\": Query,\n",
    "        \"get_current_balance\": inv_utils.get_current_balance,\n",
    "        \"next_transaction_id\": inv_utils.next_transaction_id,\n",
    "        \"user_request\": user_request or \"\",\n",
    "    }\n",
    "    SAFE_LOCALS = {\n",
    "        \"db\": db,\n",
    "        \"inventory_tbl\": inventory_tbl,\n",
    "        \"transactions_tbl\": transactions_tbl,\n",
    "    }\n",
    "\n",
    "    # Capture stdout from the executed code\n",
    "    _stdout_buf, _old_stdout = io.StringIO(), sys.stdout\n",
    "    sys.stdout = _stdout_buf\n",
    "    err_text = None\n",
    "    try:\n",
    "        exec(code, SAFE_GLOBALS, SAFE_LOCALS)\n",
    "    except Exception:\n",
    "        err_text = traceback.format_exc()\n",
    "    finally:\n",
    "        sys.stdout = _old_stdout\n",
    "    printed = _stdout_buf.getvalue().strip()\n",
    "\n",
    "    # Extract possible answers set by the generated code\n",
    "    answer = (\n",
    "        SAFE_LOCALS.get(\"answer_text\")\n",
    "        or SAFE_LOCALS.get(\"answer_rows\")\n",
    "        or SAFE_LOCALS.get(\"answer_json\")\n",
    "    )\n",
    "\n",
    "\n",
    "    return {\n",
    "        \"code\": code,            # <- ya sin etiquetas\n",
    "        \"stdout\": printed,\n",
    "        \"error\": err_text,\n",
    "        \"answer\": answer,\n",
    "        \"transactions_tbl\": transactions_tbl.all(),  # For inspection\n",
    "        \"inventory_tbl\": inventory_tbl.all(),  # For inspection\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b6bd5c3",
   "metadata": {},
   "source": [
    "You’ve checked the shelves and confirmed there’s exactly one round style under $100. Now the fun part: let’s hand the model’s plan-as-code to our executor and watch it do the work. The executor will peel out the <code><execute_python>...</execute_python></code> block, run it in a locked-down sandbox, and then show you everything that matters—what changed in the tables (before/after), any logs the plan printed, and the final, customer-friendly answer_text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26e5312b",
   "metadata": {
    "height": 215
   },
   "outputs": [],
   "source": [
    "# Execute the generated plan for the round-sunglasses question\n",
    "result = execute_generated_code(\n",
    "    full_content_round,          # the full LLM response you generated earlier\n",
    "    db=db,\n",
    "    inventory_tbl=inventory_tbl,\n",
    "    transactions_tbl=transactions_tbl,\n",
    "    user_request=prompt_round, # e.g., \"Do you have any round sunglasses in stock that are under $100?\"\n",
    ")\n",
    "\n",
    "# Peek at exactly what Python the plan executed\n",
    "utils.print_html(result[\"answer\"], title=\"Plan Execution · Extracted Answer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5364be2",
   "metadata": {},
   "source": [
    "As you can see, this is the expected result based on our previous manual analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ed6698b",
   "metadata": {},
   "source": [
    "## 2.4 Return Two Aviator Sunglasses\n",
    "\n",
    "In the previous step we only **queried** the data, so inventory and transactions were unchanged.  \n",
    "Now let’s handle a **return** scenario using the planning-in-code pattern:\n",
    "> **Request:** “Return 2 Aviator sunglasses I bought last week.”\n",
    "\n",
    "Before generating the plan, let’s **inspect the current inventory** for the *Aviator* model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40723d4b",
   "metadata": {
    "height": 232
   },
   "outputs": [],
   "source": [
    "Item = Query()                    # Create a Query object to reference fields (e.g., Item.name, Item.description)\n",
    "\n",
    "# Query: fetch all inventory rows whose 'name' is exactly \"Aviator\".\n",
    "# Notes:\n",
    "# - This is a case-sensitive equality check. \"aviator\" won't match.\n",
    "# - If you need case-insensitive matching, consider a .test(...) or .matches(...) with re.I.\n",
    "aviators = inventory_tbl.search(\n",
    "    (Item.name == \"Aviator\")\n",
    ")\n",
    "\n",
    "# Display the matched documents in a readable JSON panel\n",
    "utils.print_html(json.dumps(aviators, indent=2), title=\"Inventory status: Aviator sunglasses before return\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f4ca6a",
   "metadata": {},
   "source": [
    "Inventory confirms one Aviator SKU in stock — **SG001 (Aviator)**: **23** units at **$80** each. Now let's generate a plan to answer the prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ad9c2b",
   "metadata": {
    "height": 249
   },
   "outputs": [],
   "source": [
    "prompt_aviator = \"Return 2 Aviator sunglasses I bought last week.\"\n",
    "\n",
    "# Generate the plan-as-code (FULL content; may include <execute_python> tags)\n",
    "full_content_aviator = generate_llm_code(\n",
    "    prompt_aviator,\n",
    "    inventory_tbl=inventory_tbl,\n",
    "    transactions_tbl=transactions_tbl,\n",
    "    model=\"o4-mini\",\n",
    "    temperature=1,\n",
    ")\n",
    "\n",
    "# Inspect the LLM’s plan + code (no execution here)\n",
    "utils.print_html(full_content_aviator, title=\"Plan with Code (Full Response)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c1b2054",
   "metadata": {},
   "source": [
    "Before we execute the plan, let’s check the current status of the transactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c809d36",
   "metadata": {
    "height": 45
   },
   "outputs": [],
   "source": [
    "utils.print_html(json.dumps(transactions_tbl.all(), indent=2), title=\"Transactions Table Before Return\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb7d8bc5",
   "metadata": {},
   "source": [
    "The transaction log currently shows a single entry — the opening balance (`TXN001`) for `$500.00` recorded at `2025-10-03T09:16:59.628898`. \n",
    "\n",
    "Ready to go—execute the plan by running the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b23879c",
   "metadata": {
    "height": 215
   },
   "outputs": [],
   "source": [
    "# Execute the generated plan for the round-sunglasses question\n",
    "result = execute_generated_code(\n",
    "    full_content_aviator,          # the full LLM response you generated earlier\n",
    "    db=db,\n",
    "    inventory_tbl=inventory_tbl,\n",
    "    transactions_tbl=transactions_tbl,\n",
    "    user_request=prompt_aviator, # e.g., \"Return 2 aviator sunglasses I bought last week.\"\n",
    ")\n",
    "\n",
    "# Peek at exactly what Python the plan executed\n",
    "utils.print_html(result[\"answer\"], title=\"Plan Execution · Extracted Answer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89a3c738",
   "metadata": {},
   "source": [
    "You can see below that a new transaction has been inserted for the Aviator sunglasses return."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7900f6a",
   "metadata": {
    "height": 45
   },
   "outputs": [],
   "source": [
    "utils.print_html(json.dumps(transactions_tbl.all(), indent=2), title=\"Transactions Table After Return\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d51a0dc",
   "metadata": {},
   "source": [
    "And by running the cell below, you’ll see the Aviator stock increase to 25 (`quantity_in_stock`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa17d2a",
   "metadata": {
    "height": 147
   },
   "outputs": [],
   "source": [
    "Item = Query()                  \n",
    "\n",
    "aviators = inventory_tbl.search(\n",
    "    (Item.name == \"Aviator\")\n",
    ")\n",
    "\n",
    "utils.print_html(json.dumps(aviators, indent=2), title=\"Inventory status: Aviator sunglasses after return\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdcc858a",
   "metadata": {},
   "source": [
    "## 3. Putting It All Together: Customer Service Agent\n",
    "\n",
    "You’ve built the pieces—schema, prompt, code generator, and executor. Now let’s wire them up into a single helper that takes a natural-language request, generates a plan-as-code, executes it safely, and shows the result (plus before/after tables).\n",
    "\n",
    "**What this agent does**\n",
    "- Optionally reseeds the demo data for a clean run.\n",
    "- Generates the plan (Python inside `<execute_python>…</execute_python>`).\n",
    "- Executes the plan in a controlled namespace (TinyDB + helpers).\n",
    "- Surfaces a concise `answer_text` and renders before/after snapshots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a786b7b",
   "metadata": {
    "height": 1371
   },
   "outputs": [],
   "source": [
    "def customer_service_agent(\n",
    "    question: str,\n",
    "    *,\n",
    "    db,\n",
    "    inventory_tbl,\n",
    "    transactions_tbl,\n",
    "    model: str = \"o4-mini\",\n",
    "    temperature: float = 1.0,\n",
    "    reseed: bool = False,\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    End-to-end helper:\n",
    "      1) (Optional) reseed inventory & transactions\n",
    "      2) Generate plan-as-code from `question`\n",
    "      3) Execute in a controlled namespace\n",
    "      4) Render before/after snapshots and return artifacts\n",
    "\n",
    "    Returns:\n",
    "      {\n",
    "        \"full_content\": <raw LLM response (may include <execute_python> tags)>,\n",
    "        \"exec\": {\n",
    "            \"code\": <extracted python>,\n",
    "            \"stdout\": <plan logs>,\n",
    "            \"error\": <traceback or None>,\n",
    "            \"answer\": <answer_text/rows/json>,\n",
    "            \"inventory_after\": [...],\n",
    "            \"transactions_after\": [...]\n",
    "        }\n",
    "      }\n",
    "    \"\"\"\n",
    "    # 0) Optional reseed\n",
    "    if reseed:\n",
    "        inv_utils.create_inventory()\n",
    "        inv_utils.create_transactions()\n",
    "\n",
    "    # 1) Show the question\n",
    "    utils.print_html(question, title=\"User Question\")\n",
    "\n",
    "    # 2) Generate plan-as-code (FULL content)\n",
    "    full_content = generate_llm_code(\n",
    "        question,\n",
    "        inventory_tbl=inventory_tbl,\n",
    "        transactions_tbl=transactions_tbl,\n",
    "        model=model,\n",
    "        temperature=temperature,\n",
    "    )\n",
    "    utils.print_html(full_content, title=\"Plan with Code (Full Response)\")\n",
    "\n",
    "    # 3) Before snapshots\n",
    "    utils.print_html(json.dumps(inventory_tbl.all(), indent=2), title=\"Inventory Table · Before\")\n",
    "    utils.print_html(json.dumps(transactions_tbl.all(), indent=2), title=\"Transactions Table · Before\")\n",
    "\n",
    "    # 4) Execute\n",
    "    exec_res = execute_generated_code(\n",
    "        full_content,\n",
    "        db=db,\n",
    "        inventory_tbl=inventory_tbl,\n",
    "        transactions_tbl=transactions_tbl,\n",
    "        user_request=question,\n",
    "    )\n",
    "\n",
    "    # 5) After snapshots + final answer\n",
    "    utils.print_html(exec_res[\"answer\"], title=\"Plan Execution · Extracted Answer\")\n",
    "    utils.print_html(json.dumps(inventory_tbl.all(), indent=2), title=\"Inventory Table · After\")\n",
    "    utils.print_html(json.dumps(transactions_tbl.all(), indent=2), title=\"Transactions Table · After\")\n",
    "\n",
    "    # 6) Return artifacts\n",
    "    return {\n",
    "        \"full_content\": full_content,\n",
    "        \"exec\": {\n",
    "            \"code\": exec_res[\"code\"],\n",
    "            \"stdout\": exec_res[\"stdout\"],\n",
    "            \"error\": exec_res[\"error\"],\n",
    "            \"answer\": exec_res[\"answer\"],\n",
    "            \"inventory_after\": inventory_tbl.all(),\n",
    "            \"transactions_after\": transactions_tbl.all(),\n",
    "        },\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ddf4502",
   "metadata": {},
   "source": [
    "## 4. Try It Out (with the Customer Service Agent)\n",
    "\n",
    "Use the `customer_service_agent(...)` helper to go from a natural-language request → plan-as-code → safe execution → before/after snapshots.\n",
    "\n",
    "**Try these prompts:**\n",
    "1) **Read-only (Andrew’s example):**  \n",
    "   “Do you have any round sunglasses in stock that are under $100?”\n",
    "2) **Mutation — return:**  \n",
    "   “Return 2 Aviator sunglasses.”\n",
    "3) **Mutation — purchase:**  \n",
    "   “Purchase 3 Wayfarer sunglasses for customer Alice.”\n",
    "4) **Mutation - purchase multiple items:**\n",
    "   \"I want to buy 3 pairs of classic sunglasses and 1 pair of aviator.\"\n",
    "\n",
    "\n",
    "<div style=\"border:1px solid #93c5fd; border-left:6px solid #3b82f6; background:#eff6ff; border-radius:8px; padding:14px 16px; color:#1e3a8a; font-family:system-ui,-apple-system,Segoe UI,Roboto,Ubuntu,Cantarell,Noto Sans,sans-serif;\">\n",
    "  🔎 <strong>What does <code>reseed=True</code> do?</strong><br><br>\n",
    "  When you call <code>customer_service_agent(..., reseed=True)</code>, the agent <em>re-initializes</em> the demo data before running your prompt:\n",
    "  <ul style=\"margin:8px 0 0 18px;\">\n",
    "    <li><strong>Resets</strong> the <code>inventory_tbl</code> to the default product set.</li>\n",
    "    <li><strong>Resets</strong> the <code>transactions_tbl</code> to a single opening-balance entry.</li>\n",
    "    <li>Ensures a <strong>clean, reproducible</strong> run so results aren’t affected by previous tests.</li>\n",
    "  </ul>\n",
    "  Set <code>reseed=False</code> if you want to <strong>preserve</strong> the current state and continue from prior operations.\n",
    "</div>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16cfc8eb",
   "metadata": {
    "height": 215
   },
   "outputs": [],
   "source": [
    "prompt = \"I want to buy 3 pairs of classic sunglasses and 1 pair of aviator sunglasses.\"\n",
    "\n",
    "out = customer_service_agent(\n",
    "    prompt,\n",
    "    db=db,\n",
    "    inventory_tbl=inventory_tbl,\n",
    "    transactions_tbl=transactions_tbl,\n",
    "    model=\"o4-mini\",\n",
    "    temperature=1.0,\n",
    "    reseed=True,   # set False to keep current state of the inventory and the transactions\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39194fd8",
   "metadata": {},
   "source": [
    "## 5. Takeaways\n",
    "\n",
    "- **You let code be the plan.** Following Andrew’s “code-as-action” idea, you had the model write Python that chains the steps (filter → compute → update) and then you just ran it.\n",
    "\n",
    "- **You skipped the brittle tool soup.** Instead of piling on tiny tools or JSON plans, you used Python/TinyDB—giving the model a big, familiar toolbox that handles many query shapes with one prompt.\n",
    "\n",
    "- **You kept runs safe and visible.** You executed in a controlled namespace, captured logs/errors, and reviewed before/after tables—so you always know what changed and why."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c43bbd6f",
   "metadata": {},
   "source": [
    "<div style=\"border:1px solid #22c55e; border-left:6px solid #16a34a; background:#dcfce7; border-radius:6px; padding:14px 16px; color:#064e3b; font-family:system-ui,-apple-system,Segoe UI,Roboto,Ubuntu,Cantarell,Noto Sans,sans-serif;\">\n",
    "\n",
    "🎉 <strong>Congratulations!</strong>\n",
    "\n",
    "You just finished the lab and built an <em>agentic</em> customer service workflow. You let the model write code as the plan, ran it safely, and used simple validations to keep updates reliable. When things failed, you surfaced clear, human-readable reasons; when things worked, you saw exactly what changed via before/after snapshots.\n",
    "\n",
    "With this pattern—planning <em>in</em> code, plus transparent execution—you’re ready to design your own workflows that feel automatic, safe, and easy to extend. 🚀\n",
    "\n",
    "</div>\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
